#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""文件备份工具

功能：
1. 实现完全备份和增量备份
2. 完全备份：将源目录所有文件打包为tar.gz格式
3. 增量备份：仅备份新增或修改的文件
4. 通过MD5值比较文件变化
5. 根据日期和星期判断备份类型
"""

import os
import sys
import json
import hashlib
import tarfile
import datetime
import shutil

# 配置参数
SOURCE_DIR = '/home/data'          # 源目录
BACKUP_DIR = '/home/backup'        # 备份目录
MD5_FILE = os.path.join(BACKUP_DIR, 'data.json')  # MD5记录文件

class BackupTool:
    def __init__(self):
        # 确保目标目录存在
        os.makedirs(BACKUP_DIR, exist_ok=True)
        
        # 获取当前日期
        self.today = datetime.datetime.now()
        self.date_str = self.today.strftime('%Y%m%d')
        self.weekday = self.today.weekday()  # 0=周一, 6=周日
        
        # 设置备份文件名
        self.full_backup_file = os.path.join(BACKUP_DIR, f'data-full-{self.date_str}.tar.gz')
        self.incre_backup_file = os.path.join(BACKUP_DIR, f'data-incre-{self.date_str}.tar.gz')
    
    def calculate_md5(self, file_path, block_size=8192):
        """计算文件的MD5值，支持大文件分块处理"""
        md5_hash = hashlib.md5()
        try:
            with open(file_path, "rb") as f:
                while chunk := f.read(block_size):
                    md5_hash.update(chunk)
            return md5_hash.hexdigest()
        except Exception as e:
            print(f"计算文件 {file_path} 的MD5值失败: {str(e)}")
            return None
    
    def scan_directory(self, directory):
        """扫描目录，获取所有文件及其MD5值"""
        file_md5_dict = {}
        
        for root, _, files in os.walk(directory):
            for file in files:
                file_path = os.path.join(root, file)
                # 计算相对路径作为字典键
                rel_path = os.path.relpath(file_path, directory)
                md5_value = self.calculate_md5(file_path)
                if md5_value:
                    file_md5_dict[rel_path] = md5_value
        
        return file_md5_dict
    
    def save_md5_record(self, md5_dict):
        """保存MD5记录到文件"""
        try:
            with open(MD5_FILE, 'w', encoding='utf-8') as f:
                json.dump(md5_dict, f, ensure_ascii=False, indent=2)
            return True
        except Exception as e:
            print(f"保存MD5记录失败: {str(e)}")
            return False
    
    def load_md5_record(self):
        """从文件加载MD5记录"""
        if not os.path.exists(MD5_FILE):
            return None
        
        try:
            with open(MD5_FILE, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            print(f"加载MD5记录失败: {str(e)}")
            return None
    
    def create_tar_gz(self, source_dir, output_file, file_list=None):
        """创建tar.gz压缩包
        
        Args:
            source_dir: 源目录路径
            output_file: 输出文件名
            file_list: 要包含的文件列表（相对路径），None表示包含所有文件
        """
        try:
            with tarfile.open(output_file, 'w:gz') as tar:
                if file_list:
                    # 仅添加指定的文件
                    for rel_path in file_list:
                        abs_path = os.path.join(source_dir, rel_path)
                        if os.path.exists(abs_path):
                            tar.add(abs_path, arcname=rel_path)
                else:
                    # 添加所有文件
                    tar.add(source_dir, arcname=os.path.basename(source_dir))
            print(f"备份文件已创建: {output_file}")
            return True
        except Exception as e:
            print(f"创建备份文件失败: {str(e)}")
            if os.path.exists(output_file):
                os.remove(output_file)
            return False
    
    def full_backup(self):
        """执行完全备份"""
        print("开始执行完全备份...")
        
        # 扫描源目录，获取所有文件的MD5值
        md5_dict = self.scan_directory(SOURCE_DIR)
        
        # 创建完全备份
        if self.create_tar_gz(SOURCE_DIR, self.full_backup_file):
            # 保存MD5记录
            if self.save_md5_record(md5_dict):
                print("完全备份完成！")
                return True
        
        print("完全备份失败！")
        return False
    
    def incremental_backup(self):
        """执行增量备份"""
        print("开始执行增量备份...")
        
        # 加载上次备份的MD5记录
        old_md5_dict = self.load_md5_record()
        if not old_md5_dict:
            print("无法获取上次备份记录，切换为完全备份...")
            return self.full_backup()
        
        # 扫描当前源目录的MD5值
        current_md5_dict = self.scan_directory(SOURCE_DIR)
        
        # 找出新增或修改的文件
        changed_files = []
        for rel_path, md5_value in current_md5_dict.items():
            if rel_path not in old_md5_dict or old_md5_dict[rel_path] != md5_value:
                changed_files.append(rel_path)
        
        if not changed_files:
            print("没有发现变化的文件，无需进行增量备份")
            return True
        
        print(f"发现 {len(changed_files)} 个新增或修改的文件")
        
        # 创建增量备份
        if self.create_tar_gz(SOURCE_DIR, self.incre_backup_file, changed_files):
            # 更新MD5记录
            if self.save_md5_record(current_md5_dict):
                print("增量备份完成！")
                return True
        
        print("增量备份失败！")
        return False
    
    def run(self):
        """运行备份程序，根据条件选择备份类型"""
        # 检查源目录是否存在
        if not os.path.exists(SOURCE_DIR):
            print(f"错误：源目录 {SOURCE_DIR} 不存在！")
            return False
        
        # 检查是否需要执行完全备份
        # 条件：今天是周一 或者 data.json文件不存在
        if self.weekday == 0 or not os.path.exists(MD5_FILE):
            return self.full_backup()
        else:
            return self.incremental_backup()

if __name__ == "__main__":
    # 检查是否以管理员权限运行（Linux系统）
    if os.name == 'posix' and os.geteuid() != 0:
        print("请以管理员权限运行此脚本")
        sys.exit(1)
        
    backup_tool = BackupTool()
    success = backup_tool.run()
    
    # 根据备份结果设置退出码
    sys.exit(0 if success else 1)
    